from openai import OpenAI
from pymongo import MongoClient
from config.config import Config
from utils.regulatory_standard import get_relevant_standards, format_standards_response
import uuid
from datetime import datetime

class QASystem:
    def __init__(self):
        self.config = Config()
        self.client = MongoClient(self.config.MONGODB_URI)
        self.db = self.client[self.config.DB_NAME]
        
        self.deepseek_client = OpenAI(
            api_key=self.config.DEEPSEEK_API_KEY,
            base_url=self.config.DEEPSEEK_BASE_URL
        )
        
        from backend.document_management.embedding_processor import EmbeddingProcessor
        self.embedding_processor = EmbeddingProcessor()
    
    def semantic_search(self, query, collection_filter=None, top_k=5):
        """Find most relevant document chunks using vector similarity with collection filtering"""
        query_embedding = self.embedding_processor.generate_embedding(query)
        
        if query_embedding is None:
            return []
        
        return self.embedding_processor.get_similar_chunks(query_embedding, collection_filter, top_k)
    
    def format_answer(self, question, relevant_chunks):
        """Format answer in the required regulatory format with enhanced citations"""
        if not relevant_chunks:
            # Use regulatory_standards.py for fallback response
            relevant_standards = get_relevant_standards(question)
            standards_response = format_standards_response(relevant_standards)
            
            return {
                'regulatory_basis': 'No relevant regulation found in knowledge base',
                'original_text_link': 'No relevant original text in the knowledge base',
                'detailed_analysis': self.generate_no_result_response(question, standards_response),
                'citations': [],
                'found_in_kb': False
            }
        
        # Get best chunk and document
        best_chunk = relevant_chunks[0]['chunk']
        document = self.db.documents.find_one({'_id': best_chunk['document_id']})
        
        if not document:
            return {
                'regulatory_basis': 'Document reference error',
                'original_text_link': 'No relevant original text in the knowledge base',
                'detailed_analysis': 'Error retrieving document information.',
                'citations': [],
                'found_in_kb': False
            }
        
        # Build regulatory basis
        regulatory_basis = f"{document.get('country', 'Unknown')} - {document.get('issuing_authority', 'Unknown')} - {document.get('regulation_id', 'Unknown')}"
        
        # Generate PDF link
        if document.get('file_format') == 'pdf':
            pdf_link = f"/static/pdfs/{document['filename']}"
        else:
            pdf_link = "Document not available as PDF"
        
        # Build context from top chunks
        context = "\n\n".join([chunk['chunk']['chunk_text'] for chunk in relevant_chunks[:3]])
        
        # Build citations with document links
        citations = []
        for i, chunk_info in enumerate(relevant_chunks[:3]):
            chunk = chunk_info['chunk']
            doc_id = chunk['document_id']
            
            # Get document info for this chunk
            chunk_doc = self.db.documents.find_one({'_id': doc_id})
            if chunk_doc:
                citations.append({
                    'document_id': str(doc_id),
                    'document_name': chunk_doc.get('file_name', chunk_doc.get('filename', 'Unknown Document')),
                    'page_number': chunk.get('page_number', 1),
                    'text': f"{chunk_doc.get('file_name', chunk_doc.get('filename', 'Unknown'))} - Page {chunk.get('page_number', 1)}",
                    'excerpt': chunk['chunk_text'][:300] + "...",
                    'file_format': chunk_doc.get('file_format', chunk_doc.get('document_type', 'pdf')),
                    'similarity': chunk_info['similarity']
                })
        
        # Generate detailed analysis using DeepSeek
        analysis = self.generate_analysis(question, context, regulatory_basis, citations)
        
        return {
            'regulatory_basis': regulatory_basis,
            'original_text_link': pdf_link,
            'original_text_excerpt': best_chunk['chunk_text'][:500] + "...",
            'detailed_analysis': analysis,
            'source_document': document.get('file_name', document.get('filename', 'Unknown')),
            'confidence_score': relevant_chunks[0]['similarity'],
            'citations': citations,
            'found_in_kb': True
        }
    
    def generate_analysis(self, question, context, regulatory_basis, citations):
        """Generate detailed analysis using DeepSeek with Regulatory Affairs Specialist role"""
        
        # Build citation references for the prompt
        citation_refs = "\n".join([
            f"[{i+1}] {cite['document_name']} (Page {cite['page_number']})"
            for i, cite in enumerate(citations)
        ])
        
        system_prompt = """You are an experienced Regulatory Affairs Specialist with deep expertise in medical device regulations worldwide, including FDA (USA), NMPA (China), PMDA (Japan), ANVISA (Brazil), EU MDR, and ISO standards.

Your role is to:
1. Provide accurate, professional regulatory guidance based on official regulations
2. Use precise regulatory terminology and cite specific regulatory requirements
3. Explain compliance pathways and practical implementation steps
4. Identify potential regulatory risks and mitigation strategies
5. Reference specific sections, clauses, or requirements from regulations

Always maintain a professional, authoritative tone appropriate for regulatory professionals. Structure your responses clearly with proper citations to regulatory documents."""

        user_prompt = f"""Based on the following regulatory documentation, provide a comprehensive professional analysis as a Regulatory Affairs Specialist.

**Question:** {question}

**Regulatory Basis:** {regulatory_basis}

**Referenced Documents:**
{citation_refs}

**Relevant Regulatory Text:**
{context}

Please provide:
1. Direct answer to the regulatory question
2. Specific regulatory requirements and citations
3. Compliance pathway and practical steps
4. Any critical regulatory considerations or risks
5. Professional recommendations

Use proper regulatory terminology and cite specific requirements where applicable. Format your response in a clear, professional manner."""

        try:
            response = self.deepseek_client.chat.completions.create(
                model=self.config.DEEPSEEK_MODEL,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.3,
                max_tokens=1500
            )
            
            return response.choices[0].message.content
            
        except Exception as e:
            print(f"Error calling DeepSeek API: {e}")
            # Fallback response
            return f"""**Regulatory Analysis**

Based on the regulatory documentation in our knowledge base:

**Regulatory Reference:** {regulatory_basis}

**Key Requirements:**
{context[:800]}...

**Referenced Documents:**
{citation_refs}

*Note: Please consult the full regulatory documents for comprehensive compliance guidance. This analysis is based on available documentation and should be verified with current regulatory sources.*"""
    
    def generate_no_result_response(self, question, standards_response=""):
        """Generate response when no relevant documents are found, with regulatory standards fallback"""
        
        system_prompt = """You are an experienced Regulatory Affairs Specialist. When the knowledge base doesn't contain relevant information, you should:
1. Acknowledge that specific information is not available in the current knowledge base
2. Suggest relevant regulatory standards or documents that would typically address this question
3. Provide official sources where this information can be found (regulatory agency websites, ISO standards, etc.)
4. Give general guidance if appropriate

Be professional, helpful, and direct users to authoritative sources."""

        user_prompt = f"""The knowledge base does not contain specific information to answer this question: "{question}"

Please:
1. Acknowledge this limitation professionally
2. Suggest 2-3 relevant regulatory documents or standards that would address this topic
3. Provide official links or references where possible (e.g., FDA.gov, ISO.org, NMPA website)
4. Offer brief general guidance about the regulatory area if appropriate

Focus on being helpful while being clear about the limitation."""

        try:
            response = self.deepseek_client.chat.completions.create(
                model=self.config.DEEPSEEK_MODEL,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.3,
                max_tokens=800
            )
            
            ai_response = response.choices[0].message.content
            
            # Append the standards response if available
            if standards_response:
                return f"{ai_response}\n\n---\n\n{standards_response}"
            else:
                return ai_response
            
        except Exception as e:
            print(f"Error calling DeepSeek API for no-result response: {e}")
            # Use the standards_response if DeepSeek fails
            if standards_response:
                return f"""I apologize, but the current knowledge base does not contain specific information to answer your question about: "{question}"

{standards_response}

Please consult official regulatory sources or a qualified regulatory affairs professional for specific guidance."""
            else:
                return f"""I apologize, but the current knowledge base does not contain specific information to answer your question about: "{question}"

**Recommended Resources:**

For medical device regulatory information, please consult:

1. **FDA** (USA): https://www.fda.gov/medical-devices
   - Guidance documents, regulations, and submission requirements

2. **ISO Standards**: https://www.iso.org
   - ISO 13485 (Quality Management Systems)
   - ISO 14971 (Risk Management)
   - ISO 10993 series (Biocompatibility)

3. **NMPA** (China): https://www.nmpa.gov.cn
   - Registration requirements and technical guidelines

4. **PMDA** (Japan): https://www.pmda.go.jp/english/
   - Approval and compliance documentation

5. **Regulatory Affairs Professional Society**: https://www.raps.org
   - Regulatory resources and guidance

Please consult official regulatory sources or a qualified regulatory affairs professional for specific guidance."""
    
    def ask_question(self, question, session_id=None, user_id="default", collection_filter=None):
        """Main function to handle questions with collection filtering"""
        if not session_id:
            session_id = str(uuid.uuid4())
        
        # Perform semantic search with collection filter
        relevant_chunks = self.semantic_search(question, collection_filter)
        
        # Format answer with enhanced citations
        answer = self.format_answer(question, relevant_chunks)
        
        # Save to database
        qa_record = {
            'session_id': session_id,
            'user_id': user_id,
            'question': question,
            'question_language': 'english',
            'answer': answer,
            'collection_filter': collection_filter,
            'timestamp': datetime.now(),
            'confidence_score': answer.get('confidence_score', 0),
            'found_relevant': answer.get('found_in_kb', False)
        }
        
        self.db.qa_sessions.insert_one(qa_record)
        
        return {
            'session_id': session_id,
            'answer': answer
        }
    
    def multi_round_conversation(self, session_id, new_question):
        """Handle multi-round conversations"""
        history = list(self.db.qa_sessions.find({'session_id': session_id}).sort('timestamp', 1))
        
        messages = [
            {"role": "system", "content": "You are an experienced Regulatory Affairs Specialist for medical devices with expertise in global regulations including FDA, NMPA, PMDA, ANVISA, and ISO standards."}
        ]
        
        for record in history:
            messages.append({"role": "user", "content": record['question']})
            messages.append({"role": "assistant", "content": record['answer']['detailed_analysis']})
        
        messages.append({"role": "user", "content": new_question})
        
        try:
            response = self.deepseek_client.chat.completions.create(
                model=self.config.DEEPSEEK_MODEL,
                messages=messages,
                temperature=0.3
            )
            
            answer_text = response.choices[0].message.content
            
            qa_record = {
                'session_id': session_id,
                'user_id': 'default',
                'question': new_question,
                'answer': {'detailed_analysis': answer_text},
                'timestamp': datetime.now()
            }
            
            self.db.qa_sessions.insert_one(qa_record)
            
            return {'session_id': session_id, 'answer': answer_text}
            
        except Exception as e:
            return {'error': str(e)}